package com.sparsearr.datastructures.ztr.huffmantree;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;

/**
 * 赫夫曼编码工具类
 * 功能：字符串/文件的霍夫曼压缩与解压，基于字符出现频率构建最优前缀编码，实现数据无损压缩
 * 核心原理：
 * 1. 统计字符频率 → 构建霍夫曼树 → 生成霍夫曼编码表
 * 2. 压缩：通过编码表将原始数据转换为二进制串，按8位分组转为字节数组
 * 3. 解压：将字节数组还原为二进制串，通过编码表反向解析为原始数据
 * @Author: ZhaoTR
 * @Date: Created in 2025/8/2 21:55
 * @Description: 赫夫曼编码（无损压缩/解压工具）
 * @Version: 2.1（修复补位逻辑导致的数组越界错误）
 */
public class HuffManCode {
    // 日志对象（用于打印运行日志）
    private static final Logger logger = LogManager.getLogger(HuffManCode.class);

    public static void main(String[] args) {
        // 测试1：字符串压缩和解码（验证核心逻辑）
        testStringCompressAndDecode();

        // 测试2：文件压缩和解压（替换为实际文件路径即可使用）
        String srcFile = "test_source.txt";    // 源文件
        String zipFile = "test_compressed.huf";// 压缩文件
        String unzipFile = "test_uncompressed.txt";// 解压文件
        testFileCompressAndUncompress(srcFile, zipFile, unzipFile);
    }

    // ------------------------------ 字符串压缩/解码测试 ------------------------------
    /**
     * 测试字符串的霍夫曼压缩与解码，验证核心逻辑正确性
     */
    private static void testStringCompressAndDecode() {
        System.out.println("=== 字符串压缩解码测试 ===");
        String originalStr = "abcdefgabcdefg123456!@# 霍夫曼编码测试";
        System.out.println("原始字符串：" + originalStr);

        byte[] originalBytes = originalStr.getBytes();
        System.out.println("原始字节数组长度：" + originalBytes.length + "，内容：" + Arrays.toString(originalBytes));

        // 1. 压缩
        byte[] compressedBytes = huffManZip(originalBytes);
        System.out.println("压缩后字节数组长度：" + compressedBytes.length + "，内容：" + Arrays.toString(compressedBytes));
        System.out.println("压缩率：" + String.format("%.2f%%", (1 - (double) compressedBytes.length / originalBytes.length) * 100));

        // 2. 生成编码表（解码需使用与压缩一致的编码表）
        Map<Byte, String> huffCodes = generateHuffmanCodes(originalBytes);
        System.out.println("霍夫曼编码表：" + huffCodes);

        // 3. 解码
        Byte[] decodedBytes = decode(huffCodes, compressedBytes);
        byte[] decodedByteArray = convertByteWrapperToArray(decodedBytes);
        String decodedStr = new String(decodedByteArray);
        System.out.println("解码后字符串：" + decodedStr);

        // 验证结果
        if (originalStr.equals(decodedStr)) {
            System.out.println("字符串压缩解码成功！");
        } else {
            System.out.println("字符串压缩解码失败！");
        }
        System.out.println("----------------------------------------");
    }

    // ------------------------------ 文件压缩/解压测试 ------------------------------
    /**
     * 测试文件的霍夫曼压缩与解压，验证文件操作正确性
     * @param srcFile 源文件路径
     * @param zipFile 压缩文件路径
     * @param unzipFile 解压文件路径
     */
    private static void testFileCompressAndUncompress(String srcFile, String zipFile, String unzipFile) {
        System.out.println("=== 文件压缩解压测试 ===");
        try {
            // 若源文件不存在，创建测试文件
            File src = new File(srcFile);
            if (!src.exists()) {
                Files.write(Paths.get(srcFile), "霍夫曼编码文件压缩测试：abcdefg1234567890!@#$%^&*()_+".getBytes());
                System.out.println("已创建测试源文件：" + srcFile);
            }

            // 1. 压缩文件
            long start = System.currentTimeMillis();
            zipFile(srcFile, zipFile);
            long zipTime = System.currentTimeMillis() - start;
            System.out.println("文件压缩完成，耗时：" + zipTime + "ms");
            System.out.println("源文件大小：" + src.length() + "字节");
            System.out.println("压缩文件大小：" + new File(zipFile).length() + "字节");

            // 2. 解压文件
            start = System.currentTimeMillis();
            unzipFile(zipFile, unzipFile);
            long unzipTime = System.currentTimeMillis() - start;
            System.out.println("文件解压完成，耗时：" + unzipTime + "ms");

            // 验证源文件与解压文件一致性
            byte[] srcBytes = Files.readAllBytes(Paths.get(srcFile));
            byte[] unzipBytes = Files.readAllBytes(Paths.get(unzipFile));
            if (Arrays.equals(srcBytes, unzipBytes)) {
                System.out.println("文件压缩解压成功！");
            } else {
                System.out.println("文件压缩解压失败！");
            }
        } catch (IOException e) {
            logger.error("文件操作异常", e);
            System.err.println("文件操作失败：" + e.getMessage());
        }
        System.out.println("----------------------------------------");
    }

    // ------------------------------ 核心压缩逻辑 ------------------------------
    /**
     * 霍夫曼压缩入口方法：接收原始字节数组，返回压缩后的字节数组
     * @param originalBytes 原始字节数组
     * @return 压缩后的字节数组
     */
    private static byte[] huffManZip(byte[] originalBytes) {
        if (originalBytes == null || originalBytes.length == 0) {
            logger.warn("原始字节数组为空，无需压缩");
            return new byte[0];
        }

        try {
            // 1. 统计字节频率，生成节点列表
            List<Node1> nodes = getByteFrequencyNodes(originalBytes);
            // 2. 构建霍夫曼树
            Node1 huffmanRoot = createHuffmanTree(nodes);
            // 3. 生成霍夫曼编码表
            Map<Byte, String> huffCodes = getCodes(huffmanRoot);
            // 4. 执行压缩（编码表→二进制串→字节数组）
            return encodeBytes(originalBytes, huffCodes);
        } catch (Exception e) {
            logger.error("霍夫曼压缩失败", e);
            throw new RuntimeException("压缩失败：" + e.getMessage(), e);
        }
    }

    /**
     * 统计字节出现频率，生成霍夫曼树节点列表
     * @param bytes 原始字节数组
     * @return 包含字节+频率的节点列表
     */
    private static List<Node1> getByteFrequencyNodes(byte[] bytes) {
        List<Node1> nodes = new ArrayList<>();
        // 用Map统计每个字节的出现次数
        Map<Byte, Integer> frequencyMap = new HashMap<>();
        for (byte b : bytes) {
            frequencyMap.put(b, frequencyMap.getOrDefault(b, 0) + 1);
        }
        // 将Map转换为霍夫曼树节点（data=字节，weight=频率）
        for (Map.Entry<Byte, Integer> entry : frequencyMap.entrySet()) {
            nodes.add(new Node1(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    /**
     * 构建霍夫曼树：基于节点频率（权重）构建最优二叉树
     * @param nodes 字节频率节点列表
     * @return 霍夫曼树根节点
     */
    private static Node1 createHuffmanTree(List<Node1> nodes) {
        // 优先队列（最小堆）：确保每次取出权重最小的两个节点
        PriorityQueue<Node1> minHeap = new PriorityQueue<>(nodes);

        // 循环合并节点，直到只剩一个根节点
        while (minHeap.size() > 1) {
            // 取出权重最小的两个节点
            Node1 leftNode = minHeap.poll();  // 左子节点（权重更小）
            Node1 rightNode = minHeap.poll(); // 右子节点

            // 创建父节点（无实际字节数据，权重=左右子节点权重和）
            Node1 parentNode = new Node1(null, leftNode.weight + rightNode.weight);
            parentNode.left = leftNode;
            parentNode.right = rightNode;

            // 将父节点重新加入堆中，参与下一轮合并
            minHeap.add(parentNode);
        }

        // 最后剩余的节点即为霍夫曼树根节点
        return minHeap.poll();
    }

    /**
     * 生成霍夫曼编码表：遍历霍夫曼树，左子树记"0"，右子树记"1"
     * @param root 霍夫曼树根节点
     * @return 霍夫曼编码表（key=字节，value=二进制编码串）
     */
    private static Map<Byte, String> getCodes(Node1 root) {
        if (root == null) {
            logger.warn("霍夫曼树为空，无法生成编码表");
            return Collections.emptyMap();
        }

        Map<Byte, String> huffCodes = new HashMap<>();
        StringBuilder codeBuilder = new StringBuilder(); // 存储当前节点的编码前缀

        // 递归遍历左子树（编码加"0"）和右子树（编码加"1"）
        traverseHuffmanTree(root.left, "0", codeBuilder, huffCodes);
        traverseHuffmanTree(root.right, "1", codeBuilder, huffCodes);

        return huffCodes;
    }

    /**
     * 递归遍历霍夫曼树，生成编码表（辅助方法）
     * @param node 当前遍历节点
     * @param code 当前节点的编码位（左子树"0"，右子树"1"）
     * @param codeBuilder 编码前缀构建器
     * @param huffCodes 存储最终编码表的Map
     */
    private static void traverseHuffmanTree(Node1 node, String code, StringBuilder codeBuilder, Map<Byte, String> huffCodes) {
        // 复制编码前缀，避免递归过程中相互干扰（StringBuilder是可变对象）
        StringBuilder currentBuilder = new StringBuilder(codeBuilder);
        currentBuilder.append(code);

        if (node != null) {
            // 叶子节点（有字节数据）：保存完整编码
            if (node.data != null) {
                huffCodes.put(node.data, currentBuilder.toString());
            } else {
                // 内部节点（无数据）：继续遍历左右子树
                traverseHuffmanTree(node.left, "0", currentBuilder, huffCodes);
                traverseHuffmanTree(node.right, "1", currentBuilder, huffCodes);
            }
        }
    }

    /**
     * 执行压缩：通过编码表将原始字节数组转为压缩后的字节数组
     * 修复后的补位逻辑：
     * 1. 拼接原始编码串 → 2. 补位使编码串长度为8的倍数 → 3. 第一个字节存储补位数量（0-7）
     * 4. 后续字节存储编码串（8位一组）→ 总长度为 1 + 8n，确保无索引越界
     * @param originalBytes 原始字节数组
     * @param huffCodes 霍夫曼编码表
     * @return 压缩后的字节数组
     */
    private static byte[] encodeBytes(byte[] originalBytes, Map<Byte, String> huffCodes) {
        // 1. 将所有字节对应的编码拼接成一个长二进制串
        StringBuilder codeBuffer = new StringBuilder();
        for (byte b : originalBytes) {
            codeBuffer.append(huffCodes.get(b));
        }

        // 2. 计算补位数量：让编码串长度是8的倍数（补0到末尾）
        int totalLength = codeBuffer.length();
        int paddingBits = (8 - (totalLength % 8)) % 8; // 避免paddingBits=8（取模后为0）

        // 3. 补位：在编码串末尾添加paddingBits个0
        for (int i = 0; i < paddingBits; i++) {
            codeBuffer.append('0');
        }

        // 4. 计算压缩后字节数组长度：1（存储补位数量） + 编码串长度/8
        int byteCount = 1 + (codeBuffer.length() / 8);
        byte[] compressedBytes = new byte[byteCount];

        // 5. 第一个字节存储补位数量（0-7，直接转字节）
        compressedBytes[0] = (byte) paddingBits;

        // 6. 编码串按8位分组，存入后续字节
        int index = 1; // 从第2个字节开始存储编码数据
        for (int i = 0; i < codeBuffer.length(); i += 8) {
            // 截取8位二进制串
            String byteStr = codeBuffer.substring(i, i + 8);
            // 二进制串→字节（Integer.parseInt转为整数后强转）
            compressedBytes[index] = (byte) Integer.parseInt(byteStr, 2);
            index++;
        }

        return compressedBytes;
    }

    // ------------------------------ 核心解码逻辑 ------------------------------
    /**
     * 霍夫曼解码：通过编码表将压缩字节数组还原为原始字节数组
     * @param huffCodes 霍夫曼编码表（需与压缩时一致）
     * @param compressedBytes 压缩后的字节数组
     * @return 解码后的字节数组（包装类，便于后续转换）
     */
    private static Byte[] decode(Map<Byte, String> huffCodes, byte[] compressedBytes) {
        if (huffCodes.isEmpty() || compressedBytes == null || compressedBytes.length <= 1) {
            logger.warn("编码表为空或压缩字节数组无效，无法解码");
            return new Byte[0];
        }

        try {
            // 1. 还原编码字符串：从压缩字节数组中提取编码串（去掉补位）
            String codeStr = restoreCodeString(compressedBytes);

            // 2. 编码表反转（key=编码串，value=字节）：便于解码查找
            Map<String, Byte> reverseCodes = reverseHuffmanCodes(huffCodes);

            // 3. 遍历二进制串，匹配编码表，还原原始字节
            List<Byte> decodedList = matchHuffmanCodes(codeStr, reverseCodes);

            // 4. 转换为Byte数组返回
            return decodedList.toArray(new Byte[0]);
        } catch (Exception e) {
            logger.error("霍夫曼解码失败", e);
            throw new RuntimeException("解码失败：" + e.getMessage(), e);
        }
    }

    /**
     * 还原编码字符串：从压缩字节数组中提取编码串，去掉补位
     * @param compressedBytes 压缩字节数组
     * @return 原始编码字符串（无补位）
     */
    private static String restoreCodeString(byte[] compressedBytes) {
        StringBuilder codeBuilder = new StringBuilder();

        // 1. 第一个字节是补位数量（0-7）
        int paddingBits = compressedBytes[0] & 0xFF; // 避免负数（如byte=-1→255，这里paddingBits最大7，直接强转即可）

        // 2. 遍历后续字节（从第2个开始），转为8位二进制串
        for (int i = 1; i < compressedBytes.length; i++) {
            byte b = compressedBytes[i];
            // 字节→8位二进制串（非最后一个编码字节，必须补0到8位）
            codeBuilder.append(byteTo8BitString(b));
        }

        // 3. 去掉末尾的补位（paddingBits个0）
        if (paddingBits > 0) {
            codeBuilder.setLength(codeBuilder.length() - paddingBits);
        }

        return codeBuilder.toString();
    }

    /**
     * 字节→8位二进制串：确保返回8位（不足补前导0）
     * @param b 待转换的字节
     * @return 8位二进制串
     */
    private static String byteTo8BitString(byte b) {
        int temp = b;
        // 按位或0x100（256），确保二进制串至少9位，再截取后8位（处理负数补码）
        temp |= 0x100;
        String binaryStr = Integer.toBinaryString(temp);
        return binaryStr.substring(binaryStr.length() - 8);
    }

    /**
     * 反转霍夫曼编码表：key=编码串，value=字节（便于解码时查找）
     * @param huffCodes 原始编码表（key=字节，value=编码串）
     * @return 反转后的编码表
     */
    private static Map<String, Byte> reverseHuffmanCodes(Map<Byte, String> huffCodes) {
        Map<String, Byte> reverseCodes = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffCodes.entrySet()) {
            reverseCodes.put(entry.getValue(), entry.getKey());
        }
        return reverseCodes;
    }

    /**
     * 匹配霍夫曼编码：遍历二进制串，查找对应的字节
     * @param codeStr 原始编码字符串
     * @param reverseCodes 反转后的编码表
     * @return 解码后的字节列表
     */
    private static List<Byte> matchHuffmanCodes(String codeStr, Map<String, Byte> reverseCodes) {
        List<Byte> decodedList = new ArrayList<>();
        int currentIndex = 0; // 当前遍历位置

        while (currentIndex < codeStr.length()) {
            int matchLength = 1; // 匹配的编码长度（从1开始递增）
            Byte targetByte = null;

            // 循环查找匹配的编码（直到找到对应的字节）
            while (currentIndex + matchLength <= codeStr.length()) {
                String currentCode = codeStr.substring(currentIndex, currentIndex + matchLength);
                targetByte = reverseCodes.get(currentCode);
                if (targetByte != null) {
                    break; // 找到匹配的编码，退出循环
                }
                matchLength++;
            }

            // 未找到匹配编码，抛出异常
            if (targetByte == null) {
                throw new RuntimeException("未找到匹配的霍夫曼编码：" + codeStr.substring(currentIndex, Math.min(currentIndex + 10, codeStr.length())) + "...");
            }

            // 找到字节，加入列表，更新遍历位置
            decodedList.add(targetByte);
            currentIndex += matchLength;
        }

        return decodedList;
    }

    // ------------------------------ 文件操作方法 ------------------------------
    /**
     * 文件压缩：将源文件通过霍夫曼编码压缩为目标文件
     * 注：压缩文件包含两部分数据：压缩字节数组 + 霍夫曼编码表（用于解压）
     * @param srcFile 源文件路径
     * @param dstFile 压缩文件路径
     */
    public static void zipFile(String srcFile, String dstFile) {
        FileInputStream fis = null;
        FileOutputStream fos = null;
        ObjectOutputStream oos = null;

        try {
            // 1. 读取源文件字节数据
            fis = new FileInputStream(srcFile);
            byte[] originalBytes = new byte[fis.available()];
            fis.read(originalBytes);

            // 2. 执行霍夫曼压缩
            byte[] compressedBytes = huffManZip(originalBytes);

            // 3. 生成编码表（需与压缩字节一起存储，用于解压）
            List<Node1> nodes = getByteFrequencyNodes(originalBytes);
            Node1 huffmanRoot = createHuffmanTree(nodes);
            Map<Byte, String> huffCodes = getCodes(huffmanRoot);

            // 4. 写入压缩文件（序列化存储字节数组和编码表）
            fos = new FileOutputStream(dstFile);
            oos = new ObjectOutputStream(fos);
            oos.writeObject(compressedBytes); // 写入压缩字节数组
            oos.writeObject(huffCodes);       // 写入霍夫曼编码表

            logger.info("文件压缩成功：{} → {}", srcFile, dstFile);
        } catch (IOException e) {
            logger.error("文件压缩失败", e);
            throw new RuntimeException("文件压缩失败：" + e.getMessage(), e);
        } finally {
            // 关闭流（按创建顺序反向关闭，避免资源泄露）
            closeStream(oos, fos, fis);
        }
    }

    /**
     * 文件解压：将霍夫曼压缩文件解压为目标文件
     * @param srcFile 压缩文件路径
     * @param dstFile 解压文件路径
     */
    public static void unzipFile(String srcFile, String dstFile) {
        FileInputStream fis = null;
        ObjectInputStream ois = null;
        FileOutputStream fos = null;

        try {
            // 1. 读取压缩文件数据（字节数组 + 编码表）
            fis = new FileInputStream(srcFile);
            ois = new ObjectInputStream(fis);
            byte[] compressedBytes = (byte[]) ois.readObject(); // 读取压缩字节数组
            Map<Byte, String> huffCodes = (Map<Byte, String>) ois.readObject(); // 读取编码表

            // 2. 执行霍夫曼解码
            Byte[] decodedBytes = decode(huffCodes, compressedBytes);
            byte[] dstBytes = convertByteWrapperToArray(decodedBytes);

            // 3. 写入解压文件
            fos = new FileOutputStream(dstFile);
            fos.write(dstBytes);

            logger.info("文件解压成功：{} → {}", srcFile, dstFile);
        } catch (IOException | ClassNotFoundException e) {
            logger.error("文件解压失败", e);
            throw new RuntimeException("文件解压失败：" + e.getMessage(), e);
        } finally {
            // 关闭流
            closeStream(ois, fis, fos);
        }
    }

    // ------------------------------ 工具辅助方法 ------------------------------
    /**
     * Byte包装类数组 → 原始byte数组
     * @param wrapperArray Byte包装类数组
     * @return 原始byte数组
     */
    private static byte[] convertByteWrapperToArray(Byte[] wrapperArray) {
        if (wrapperArray == null || wrapperArray.length == 0) {
            return new byte[0];
        }
        byte[] primitiveArray = new byte[wrapperArray.length];
        for (int i = 0; i < wrapperArray.length; i++) {
            primitiveArray[i] = wrapperArray[i];
        }
        return primitiveArray;
    }

    /**
     * 生成霍夫曼编码表（单独提取，便于测试）
     * @param originalBytes 原始字节数组
     * @return 霍夫曼编码表
     */
    private static Map<Byte, String> generateHuffmanCodes(byte[] originalBytes) {
        List<Node1> nodes = getByteFrequencyNodes(originalBytes);
        Node1 root = createHuffmanTree(nodes);
        return getCodes(root);
    }

    /**
     * 关闭多个流（通用方法，避免重复代码）
     * @param streams 待关闭的流（可传入多个）
     */
    private static void closeStream(Closeable... streams) {
        for (Closeable stream : streams) {
            if (stream != null) {
                try {
                    stream.close();
                } catch (IOException e) {
                    logger.error("流关闭失败", e);
                }
            }
        }
    }
}

/**
 * 霍夫曼树节点类：存储字节数据和对应的权重（出现频率）
 * 实现Comparable接口，用于优先队列的排序（按权重升序）
 */
class Node1 implements Comparable<Node1> {
    Byte data;       // 存储的字节数据（叶子节点有效，内部节点为null）
    int weight;      // 权重（字节出现频率）
    Node1 left;      // 左子节点（编码"0"）
    Node1 right;     // 右子节点（编码"1"）

    /**
     * 构造方法
     * @param data 字节数据
     * @param weight 权重（出现频率）
     */
    public Node1(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    /**
     * 比较器：按权重升序排序（用于优先队列）
     * @param o 另一个节点
     * @return 负数：当前节点权重小；0：权重相等；正数：当前节点权重大
     */
    @Override
    public int compareTo(Node1 o) {
        return this.weight - o.weight;
    }

    /**
     * 节点字符串表示（便于调试）
     */
    @Override
    public String toString() {
        return "Node1{data=" + data + ", weight=" + weight + "}";
    }

    /**
     * 前序遍历（用于调试霍夫曼树结构）
     */
    public void preOrder() {
        System.out.println(this);
        if (left != null) {
            left.preOrder();
        }
        if (right != null) {
            right.preOrder();
        }
    }
}